﻿using System;
using System.Collections.Generic;
using System.Drawing;
using System.Linq;
using System.Text;
using Emgu.CV;
using Emgu.CV.CvEnum;
using Emgu.CV.Structure;

namespace ImageAlignment
{
  class ForwardAdditiveAlgorithm
  {
    public static void Run(Image<Gray, Byte> templateImage, Rectangle templateRect, Image<Gray, Byte> image)
    {
    	// Some constants for iterative minimization process.
	    const float eps = 1E-5f; // Threshold value for termination criteria.
	    const int maxIter = 100;  // Maximum iteration count.
	
	    // Here we will store internally used images.
      // Gradient of I in X direction.
	    var gradIx = new Image<Gray, short>(image.Width, image.Height);	   
      // Gradient of I in Y direction.
	    var gradIy = new Image<Gray, short>(image.Width, image.Height);

      // Here we will store matrices.
      var W = new Matrix<Single>(3, 3);  // Current value of warp W(x,p)
      var X = new Matrix<Single>(3, 1);  // Point in coordinate frame of T.
      var Z = new Matrix<Single>(3, 1); ;  // Point in coordinate frame of I.

      var H = new Matrix<Single>(3, 3);  // Hessian
      var iH = new Matrix<Single>(3, 3); // Inverse of Hessian
      var b = new Matrix<Single>(3, 1);  // Vector in the right side of the system of linear equations.
      var delta_p = new Matrix<Single>(3, 1); // Parameter update value.
      
      // Get current time. We will use it later to obtain total calculation time.
      var start_time = DateTime.Now;

      /*
       *  Precomputation stage.
       */

      // Calculate gradient of I.
      CvInvoke.cvSobel(image.Ptr, gradIx.Ptr, 1, 0, 3);
//      CvInvoke.cvConvertScale(gradIx.Ptr, gradIx.Ptr, 0.125, 0); // Normalize
      CvInvoke.cvSobel(image.Ptr, gradIy.Ptr, 0, 1, 3);
      CvInvoke.cvConvertScale(gradIy.Ptr, gradIy.Ptr, 0.125, 0); // Normalize

      var imageGradX = AuxFunction.GetArray(gradIx);
      var imageGradY = AuxFunction.GetArray(gradIy);

      var imageArray = AuxFunction.GetArray(image);

      /*
       *  Iteration stage.
       */

      // Here we will store parameter approximation. 
      float wz_a=0, tx_a=0, ty_a=0;

      // Here we will store current mean error value.
      float mean_error=0;

      // Iterate
      int iter=0; // number of current iteration
      while(iter < maxIter)
      {
	      iter++; // Increment iteration counter

	      mean_error = 0; // Set mean error value with zero

	      int pixel_count=0; // Count of processed pixels
      	
	      AuxFunction.InitWarp(W, wz_a, tx_a, ty_a); // Init warp W(x, p)
	      AuxFunction.SetAll(H, 0); // Set Hessian with zeroes
	      AuxFunction.SetAll(b, 0); // Set b matrix with zeroes
        
	      // (u,v) - pixel coordinates in the coordinate frame of T.
	      int u, v;

	      // Walk through pixels in the template T.
	      int i, j;
	      for(i=0; i<templateRect.Width; i++)
	      {
		      u = i + templateRect.X;

		      for(j=0; j<templateRect.Height; j++)
		      {
			      v = j + templateRect.Y;

			      // Set vector X with pixel coordinates (u,v,1)
		        X = AuxFunction.CreateVector(u, v, 1);

			      // Warp Z=W*X
			      Z = W.Mul(X);

			      // (u2,v2) - pixel coordinates in the coordinate frame of I.
		        var u2 = Z.Data[0, 0];
		        var v2 = Z.Data[1, 0];

			      // Get the nearest integer pixel coords (u2i;v2i).
			      var u2i = (int)u2;
			      var v2i = (int)v2;

            // check if pixel is inside I.
			      if (u2i >= 0 && u2i < image.Width && v2i >= 0 && v2i < image.Height)
			      {
				      pixel_count++;

				      // Evaluate gradient of I at W(x,p) with subpixel accuracy
				      // using bilinear interpolation.
			        float Ix = AuxFunction.Interpolate(imageGradX, u2, v2);
              float Iy = AuxFunction.Interpolate(imageGradY, u2, v2);
              
				      // Calculate steepest descent image's element.
				      var stdesc = new float[3]; // an element of steepest descent image
				      stdesc[0] = (float)(-v*Ix+u*Iy);
				      stdesc[1] = (float)Ix;
				      stdesc[2] = (float)Iy;

				      // Calculate intensity of a transformed pixel with sub-pixel accuracy
				      // using bilinear interpolation.

			        float I2 = AuxFunction.Interpolate(imageArray, u2, v2);
				      
				      // Calculate image difference D = T(x)-I(W(x,p)).
				      float D = (float)templateImage[v, u].Intensity - I2;
      				
				      // Update mean error value.
				      mean_error += Math.Abs(D);

				      // Add a term to b matrix.
			        b[0, 0] += stdesc[0]*D;
              b[1, 0] += stdesc[1]*D;
              b[2, 0] += stdesc[2]*D;

				      // Add a term to Hessian.
				      int l,m;
				      for(l=0;l<3;l++)
				      {
					      for(m=0;m<3;m++)
					      {
						      H[l, m] += stdesc[l]*stdesc[m];
					      }
				      }
			      }	 
		      }
	      }

	      // Finally, calculate mean error.
	      if(pixel_count!=0)
		      mean_error /= pixel_count;

	      // Invert Hessian.
	      double inv_res = CvInvoke.cvInvert(H.Ptr, iH.Ptr, INVERT_METHOD.CV_LU);
	      if (inv_res == 0)
	      {
		      Console.WriteLine("Error: Hessian is singular.\n");
		      return;
	      }

	      // Find parameter increment. 
	      delta_p = iH.Mul(b);
	      float delta_wz = delta_p[0, 0];
	      float delta_tx = delta_p[1, 0];
	      float delta_ty = delta_p[2, 0];

	      // Update parameter vector approximation.
	      wz_a += delta_wz;
	      tx_a += delta_tx;
	      ty_a += delta_ty;

	      // Print diagnostic information to screen.
	      Console.WriteLine("iter={0} dwz={1} dtx={2} dty={3} mean_error={4}\n", 
		      iter, delta_wz, delta_tx, delta_ty, mean_error);

	      // Check termination critera.
	      if(Math.Abs(delta_wz) < eps && Math.Abs(delta_tx) < eps && Math.Abs(delta_ty) < eps) break;
      }

      // Get current time and obtain total time of calculation.
      var finish_time = DateTime.Now;
      double total_time = (finish_time - start_time).TotalSeconds;

      // Print summary.
      Console.WriteLine("===============================================\n");
      Console.WriteLine("Algorithm: forward additive.\n");
      Console.WriteLine("Caclulation time: {0} sec.\n", total_time);
      Console.WriteLine("Iteration count: {0}\n", iter);
      Console.WriteLine("Approximation: wz_a={0} tx_a={1} ty_a={2}\n", wz_a, tx_a, ty_a);
      Console.WriteLine("Epsilon: {0}\n", eps);
      Console.WriteLine("Resulting mean error: {0}\n", mean_error);
      Console.WriteLine("===============================================\n");

      // Show result of image alignment.
      AuxFunction.InitWarp(W, wz_a, tx_a, ty_a);
      AuxFunction.DrawWarpedRectangle(image, templateRect, W);
      CvInvoke.cvSetImageROI(templateImage, templateRect);
      CvInvoke.cvShowImage("template", templateImage);
      CvInvoke.cvShowImage("image", image);
      CvInvoke.cvResetImageROI(templateImage);
    }
  }
}
